import tensorflow as tf
import numpy as np
import tensorflow.keras.backend as K


'''
实现正态分布的基本功能：
参数初始化、样本采集、计算给定值的概率和对数概率、
'''
class Normal:
    '''
    mean: 表示正态分布的均值，是一个 TensorFlow 张量
    logstd: 表示正态分布的对数标准差
    std: 正态分布的标准差，从 logstd 通过指数计算得出
    group_ndims: 表示从最后面的维度开始计算的分组维度数，用于一起计算这些维度的概率
    dtype: 表示分布中样本的类型
    '''

    def __init__(self, mean, logstd=None, group_ndims=0):
        '''
        对 mean 和 logstd 进行类型转换和有效性检查，确保它们可以广播到相同的形状
        '''
        self._mean = tf.convert_to_tensor(mean)

        if logstd is None:
            raise ValueError('Log standard deviation must not be None')
        self._logstd = tf.convert_to_tensor(logstd)

        std = tf.exp(self._logstd)
        std = tf.debugging.check_numerics(std, "exp(logstd)")
        self._std = std
        self._group_ndims = group_ndims
        self._dtype = self._mean.dtype
        try:
            tf.broadcast_static_shape(self._mean.get_shape(), self._std.get_shape())
        except ValueError:
            raise ValueError(
                "mean and logstd should be boardcastable to match each other. "
                "({} vs {})".format(self._mean.get_shape(), self._logstd.get_shape())
            )

    @property
    def dtype(self):
        """
        The sample type of the distribution
        """
        return self._dtype

    @property
    def mean(self):
        """
        The mean of the Normal distribution
        """
        return self._mean

    @property
    def logstd(self):
        """
        The log standard deviation of the Normal distribution
        """
        return self._logstd

    @property
    def std(self):
        """
        The standard deviation of the Normal distribution
        """
        return self._std

    @property
    def group_ndims(self):
        """
        The number of dimension counted from the end, that are grouped into a single
        event, so that their probabilities are calculated together
        """
        return self._group_ndims

    def get_batch_shape(self):
        # 计算并返回分布的批量形状，确保 mean 和 logstd 是可以广播的
        try:
            tf.broadcast_static_shape(self._mean.get_shape(), self._std.get_shape())
        except ValueError:
            raise ValueError(
                "mean and logstd should be boardcastable to match each other. "
                "({} vs {})".format(self._mean.get_shape(), self._logstd.get_shape())
            )
        return self.mean.get_shape()

    @property
    def batch_shape(self):
        static_batch_shape = self.get_batch_shape()
        return static_batch_shape

    def _sample(self, n_samples):
        # 生成样本
        batch_shape = tf.shape(self.mean)[0]
        hidden_dim = K.int_shape(self.mean)[1:]

        mean, std = self.mean, self.std
        shape = (n_samples, batch_shape) + hidden_dim
        samples = tf.random.normal(shape=shape, dtype=self.dtype) * std + mean
        return samples

    def sample(self, n_samples=1):
        # 生成给定数量的样本
        if n_samples is None or n_samples == 1:
            samples = self._sample(n_samples=1)
            samples = tf.squeeze(samples, axis=0)
        else:
            samples = self._sample(n_samples)

        return samples

    def _log_prob(self, given):
        c = -0.5 * np.log(2 * np.pi)
        precision = tf.exp(-2 * self.logstd)
        precision = tf.debugging.check_numerics(precision, "precision")
        return c - self.logstd - 0.5 * precision * tf.square(given - self.mean)

    def _prob(self, given):
        return tf.exp(self._log_prob(given))

    def _check_input_shape(self, given):
        # 检查输入值的形状是否能够与分布的原批量形状和值的形状相匹配，确保可以进行广播
        given = tf.convert_to_tensor(given, dtype=self.dtype)

        err_msg = "The given argument should be able to broadcast to" \
                  "match batch_shape + value_shape of the distribution"

        if (given.get_shape() and self.get_batch_shape() and self._get_value_shape()):


            static_sample_shape = tf.TensorShape(
                self.get_batch_shape().as_list() +
                self._get_value_shape().as_list()
            )
            try:
                tf.broadcast_static_shape(given.get_shape(),
                                          static_sample_shape)
            except ValueError:
                raise ValueError(
                    err_msg + " ({} vs. {} + {})".format(
                        given.get_shape(), self.get_batch_shape(),
                        self._get_value_shape())
                )
        return given

    def _get_value_shape(self):
        # 返回样本的值形状，这里返回的是一个空的张量形状（表示标量）。
        return tf.TensorShape([])

    def log_prob(self, given, use_group_ndims=True):
        # 计算给定值的对数概率密度
        given = self._check_input_shape(given)
        log_p = self._log_prob(given)

        if use_group_ndims:
            log_p = tf.reduce_sum(log_p, tf.range(-self._group_ndims, 0))
        else:
            log_p = tf.reduce_sum(log_p, [])

        return log_p

    def prob(self, given, use_group_ndims=True):
        # 计算给定值的概率密度
        given = self._check_input_shape(given)
        p = self._prob(given)

        if use_group_ndims:
            p = tf.reduce_prod(p, tf.range(-self._group_ndims, 0))
        else:
            p = tf.reduce_prod(p, [])
        return p